package com.gooluke.mds.common.util;

import com.gooluke.mds.datasource.DatasourceContextHolder;
import com.gooluke.mds.datasource.TransactionData;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;
import org.springframework.transaction.PlatformTransactionManager;
import org.springframework.transaction.TransactionDefinition;
import org.springframework.transaction.TransactionStatus;
import org.springframework.transaction.support.DefaultTransactionDefinition;

import java.util.ArrayList;
import java.util.List;

/**
 * @author gooluke
 * description 编程式事务，多数据源同时操作，提交、回滚
 * datetime 2025-05-24 4:28
 */
@Component
@Slf4j
public class TransactionUtil {

    @Autowired
    private PlatformTransactionManager transactionManager;

    public void execute(List<String> dsNames, Runnable runnable) {
        execute(dsNames, runnable, true);
    }

    public void execute(List<String> dsNames, Runnable runnable, boolean throwWhenException) {
        String originalDataSource = DatasourceContextHolder.getDatasource();
        List<TransactionData> statusList = new ArrayList<>();
        Exception exception = null;
        try {
            for (String dsName : dsNames) {
                DatasourceContextHolder.setDatasource(dsName);
                DefaultTransactionDefinition def = new DefaultTransactionDefinition();
                def.setPropagationBehavior(TransactionDefinition.PROPAGATION_REQUIRES_NEW);
                TransactionStatus status = transactionManager.getTransaction(def);
                statusList.add(new TransactionData(dsName, status));
                runnable.run();
            }
        } catch (Exception e) {
            //先不抛出，等回滚完成再抛出异常
            exception = e;
        } finally {
            // 逆序提交/回滚事务
            for (int i = statusList.size() - 1; i >= 0; i--) {
                TransactionData transactionData = statusList.get(i);
                log.info("Switching to datasource: {}", transactionData.getDataSourceName());
                DatasourceContextHolder.setDatasource(transactionData.getDataSourceName());
                TransactionStatus status = transactionData.getTransactionStatus();
                if (!status.isCompleted()) {  // 检查事务是否未完成
                    if (exception == null) {
                        transactionManager.commit(status);
                    } else {
                        transactionManager.rollback(status);
                    }
                }
            }
            // 如果最开始有设置过数据源，则恢复原始数据源
            if (originalDataSource != null) {
                DatasourceContextHolder.setDatasource(originalDataSource);
            }
        }
        if (exception != null && throwWhenException) {
            throw new RuntimeException(exception);
        }
    }

}
